import os,sys,re
import torch

currdir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(currdir)
sys.path.append(currdir + '/../../PytorchToCaffe')
print(sys.path)


from pfld.pfld_resnet import PFLDResNet as PFLDInference


input_img_size = 112  # define input size

device = 'cpu'

model_path = sys.argv[1]

checkpoint = torch.load(model_path, map_location=torch.device(device))
net = PFLDInference()
net.load_state_dict(checkpoint)
net.eval()
net.to(device)


netname = os.path.splitext(os.path.split(model_path)[1])[0]
import pytorch_to_caffe
dummy_input = torch.randn(1, 3, 112, 112).to(device)
pytorch_to_caffe.trans_net(net, dummy_input, netname)
pytorch_to_caffe.save_prototxt('{}.prototxt'.format(netname))
pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(netname))


